from copy import copy

import cv2
import numpy as np


def pre_canny(img, eq_hist=True, canny_th1=150, canny_th2=255):
    """
    预处理 - canny
    bgr_img -> bw_edge (black white img containing edges)
    
    步骤:
    bgr -> 灰度图 -> 直方图均衡化(可选) -> canny边缘检测
    
    :param img: raw bgr image
    :param eq_hist: whether us histogram equalization
    :param canny_th1: parameter for cv2.Canny()
    :param canny_th2: parameter for cv2.Canny()
    """
    # step -- convert to gray and median blur
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    gray = cv2.medianBlur(gray, ksize=3)
    # blur_img = cv2.GaussianBlur(gray, (5, 5), 0, 0) #高斯滤波
    
    # step -- equalize gray
    if eq_hist:
        gray = cv2.equalizeHist(gray) # 直方图均衡化处理
    
    # step -- detect edge using canny
    bw_edge = cv2.Canny(gray, threshold1=canny_th1,threshold2=canny_th2 ) # 边缘检测
    
    return bw_edge


def lane_by_hough(bw_edge, seg_y=300, rho=2, theta=np.pi*4/180, th_votes=120, 
        min_theta=np.pi*5/180, max_theta=np.pi*60/180):
    """
    用霍夫直线算法检测出车道线
    bw_edge -> lane
    
    :param bw_edge: 边缘检测处理输出的二值图
    :param seg_y: 分割图像上半部分与下半部分的水平直线的y坐标. 图像的上半部分通常是
        场外景象, 存在大量干扰, 因此将上半部分图像置零(涂黑), 只处理图像下半部以减少
        场外干扰, 甚至降低霍夫算法的运算量.
    :param rho: 传给 cv2.HoughLines() 的参数
    :param theta: 传给 cv2.HoughLines() 的参数
    :param th_votes: cv2.HoughLines() 的threshold参数
    :param min_theta: 传给 cv2.HoughLines() 的参数, 用于计算左侧车道线, 右侧车道线的min_theta由对称性算出
    :param max_theta: 同 min_theta
    
    :returns: line_left, line_right, 若检测到车道线, line 以[k, b]的格式给出; 若没有检测
        到车道则line的值为None
    """
    # step -- 场外区域置零
    # 出于性能上的考虑(也为了避免麻烦), 不检查 out_y 数值是否越界, 如果越界的话, 会由numpy报错
    bw_edge[0:seg_y-1] = 0
    
    # step -- 霍夫直线检测
    # rho, theta: resolation in pixels and radian, 值越小, 分辨率越高, 内存占用越多(几kB), 
    #   theta 越大, 运算量越大, rho 与运算量无关
    # min_theta, max_theta: radian, from 0 to PI, and min_theta <= max_theta
    min_theta_right = np.pi - max_theta
    max_theta_right = np.pi - min_theta
    lines_left = cv2.HoughLines(bw_edge, rho=rho, theta=theta, threshold=th_votes, min_theta=min_theta, max_theta=max_theta)  # angle for left lane
    lines_right = cv2.HoughLines(bw_edge, rho=rho, theta=theta, threshold=th_votes, min_theta=min_theta_right, max_theta=max_theta_right)  # angle for right lane
    
    # step -- 筛选车道线
    # 利用图像下半部分某条水平线与霍夫直线相交, 找到交点最内侧的直线作为车道线
    bw_height = bw_edge.shape[1]
    test_y = bw_height - 20
    
    # 把直线的参数从 rho, theta 转为 k, b
    if lines_left is not None:
        for l in lines_left:
            rho, theta = l[0]
            x0, y0 = rho*np.cos(theta), rho*np.sin(theta)
            
            k = np.tan(np.pi/2 + theta)
            b = y0 - k*x0
            
            l[0][0], l[0][1] = k, b
    
    if lines_right is not None:
        for l in lines_right:
            rho, theta = l[0]
            x0, y0 = rho*np.cos(theta), rho*np.sin(theta)
            
            k = np.tan(np.pi/2 + theta)
            b = y0 - k*x0
            
            l[0][0], l[0][1] = k, b
    
    # test_y = k*x + b --> x = (test_y - b) / k
    line_left =  None if lines_left is None else max(lines_left, key=lambda l: (test_y - l[0][1]) / l[0][0]) [0]
    line_right = None if lines_right is None else min(lines_right, key=lambda l: (test_y - l[0][1]) / l[0][0]) [0]
    
    # return lines_left, lines_right  # debug
    return line_left, line_right


def lane_coordinate(lane_left, lane_right, center):
    """
    根据车道线生成坐标
    lanes -> coordinates
    
    原理:
    以 center 为原点, 做水平线, 与左右车道线相交. 以交点与原点的距离作为坐标.
    相当于将车道线的参数从 [k, b] 转换为 [k, x], 从而便于机器人使用.
    
    :param lane_left: [k, b]
    :param lane_right:
    :param center: 用于计算车道线坐标的原点
    """
    lane_left, lane_right = copy(lane_left), copy(lane_right)  # avoid accidental change
    cx, cy = center
    
    if lane_left is not None:
        x1 = (cy - lane_left[1]) / lane_left[0]
        lane_left[1] = x1 - cx
    else:
        dx1 = None
    
    if lane_right is not None:
        x2 = (cy - lane_right[1]) / lane_right[0]
        lane_right[1] = x2 - cx
    else:
        dx2 = None
    
    return lane_left, lane_right
    

def draw_hough_lines(img, lines, color, thickness=1):
    if lines is None:
        return img
    
    # calculated by rho, theta
    # for l in lines:
        # rho, theta = l[0]
        # a = np.cos(theta)
        # b = np.sin(theta)
        # x0 = a * rho
        # y0 = b * rho
        # x1 = int(x0 + 1000 * (-b))
        # y1 = int(y0 + 1000 * (a))
        # x2 = int(x0 - 1000 * (-b))
        # y2 = int(y0 - 1000 * (a))
    
    # calculated by k, b
    for l in lines:
        k, b = l[0]
        x1, x2 = -1000, 1000
        y1, y2 = int(k*x1+b), int(k*x2+b)
        cv2.line(img, (x1, y1), (x2, y2), color, thickness)
    
    return img


# def test_hough_img(finame=None, raw_img=None):
    # if raw_img is None:
        # raw_img = cv2.imread(finame)
    
    # bw = pre_canny(raw_img, eq_hist=True)
        
    # seg_y = 200  # roi window
    # lleft, lright = lane_by_hough(bw, seg_y=seg_y)
    # draw_hough_lines(raw_img, lleft, [255, 0, 0])
    # draw_hough_lines(raw_img, lright, [0, 0, 255])
    
    # cv2.line(raw_img, (0, seg_y-1), (raw_img.shape[1], seg_y-1), [0,255,255], 2)  # draw roi horizontal line
    # cv2.imshow('raw', raw_img)
    # cv2.imshow('bw', bw)


# def test_hough_video(cap):
    # ret, raw_img = cap.read()
    # if ret:
        # test_hough_img(raw_img=raw_img)


def test_hough_lane_video(cap):
    ret, raw_img = cap.read()
    if ret:        
        bw = pre_canny(raw_img, eq_hist=True)
            
        seg_y = 200
        cx, cy = 320, 240
        lane_left, lane_right = lane_by_hough(bw, seg_y=seg_y)
        lleft, lright = lane_coordinate(copy(lane_left), copy(lane_right), center=(cx, cy))
    
        # draw(raw_img, bw, lane_left, lane_right, lleft, lright, (cx, cy))


def draw(img, bw, lane_left_kb, lane_right_kb, lane_left_kx, lane_right_kx, control_center):
    cx, cy = control_center
    
    font = cv2.FONT_HERSHEY_COMPLEX_SMALL
    font_scale = 2

    # draw lane
    0 if lane_left_kb  is None else draw_hough_lines(img, [[lane_left_kb]], [255, 0, 0])
    0 if lane_right_kb is None else draw_hough_lines(img, [[lane_right_kb]], [0, 0, 255])
    
    # draw insect
    cv2.circle(img, (cx, cy), 5, [0,0,255], -1)
    if lane_left_kx is not None:
        dx1 = lane_left_kx[1]
        cv2.circle(img, (int(cx+dx1), cy), 3, [0,0,255], -1)
        cv2.putText(img, '%.1f' % dx1, (int(cx+dx1), cy+10), font, font_scale, [0, 255, 255])
    if lane_right_kx is not None:
        dx2 = lane_right_kx[1]
        cv2.circle(img, (int(cx+dx2), cy), 3, [0,0,255], -1)
        cv2.putText(img, '%.1f' % dx2, (int(cx+dx2), cy+10), font, font_scale, [0, 255, 255])
    
    # cv2.line(img, (0, seg_y-1), (raw_img.shape[1], seg_y-1), [0,255,255], 2)  # draw roi horizontal line
    cv2.imshow('raw', img)
    cv2.imshow('bw', bw)
    cv2.waitKey(1)
    

if __name__ == '__main__':
    import os
    from os import pardir
    from benchmark import Benchmark
    
    # ===== directory part =====
    dir_ = os.path.dirname(__file__)
    dir_img = os.path.join(dir_, r'../../../img_and_videos')
    dir_img = os.path.abspath(dir_img)
    
    # ===== img test =====
    # pth_img = os.path.join(dir_img, 'lane-320x240.jpg')
    # test_hough_img(pth_img)
    
    # ====== video test ======
    pth_vid = os.path.join(dir_img, 'lane-to-morning.mp4')
    camera = cv2.VideoCapture(pth_vid)
    
    bm = Benchmark()
    while True:
        # test_hough_video(camera)
        bm.begin()
        test_hough_lane_video(camera)
        bm.end()
        
        if bm.n % 20 == 0:
            bm.print_stats()
    
        if cv2.waitKey(50) == 27:  # Esc
            break




